use clap::Parser;
use std::fs::{self, File};
use std::path::Path;
use std::io::{self, Write};
use std::fmt::Write as FmtWrite;

#[derive(Parser)]
#[command(name = "tree")]
#[command(about = "类似树状结构的目录和文件列表显示工具")]
struct Args {
    /// 要显示的根目录
    #[arg(default_value = ".")]
    path: String,

    /// 显示所有文件，包括隐藏文件
    #[arg(short = 'a', long = "all")]
    all: bool,

    /// 仅显示目录
    #[arg(short = 'd', long = "dirs-only")]
    dirs_only: bool,

    /// 限制目录遍历的深度
    #[arg(short = 'L', long = "level")]
    level: Option<usize>,

    /// 不显示文件/目录统计数量
    #[arg(short = 'n', long = "no-count")]
    no_count: bool,

    /// 输出到文件而不是标准输出
    #[arg(short = 'o', long = "output")]
    output: Option<String>,

    /// 显示文件大小
    #[arg(short = 's', long = "size")]
    show_size: bool,
}

#[derive(Default)]
struct Stats {
    dirs: usize,
    files: usize,
}

struct TreeConfig {
    show_all: bool,
    dirs_only: bool,
    max_depth: Option<usize>,
    show_count: bool,
    show_size: bool,
}

fn main() {
    let args = Args::parse();
    
    let config = TreeConfig {
        show_all: args.all,
        dirs_only: args.dirs_only,
        max_depth: args.level,
        show_count: !args.no_count,
        show_size: args.show_size,
    };
    
    let path = Path::new(&args.path);
    if !path.exists() {
        eprintln!("错误: 路径 '{}' 不存在", args.path);
        std::process::exit(1);
    }
    
    let mut output = String::new();
    
    // 将根目录添加到输出
    writeln!(output, "{}", path.display()).unwrap();
    
    let mut stats = Stats::default();
    
    if let Err(e) = traverse_directory(path, &config, "", true, 0, &mut stats, &mut output) {
        eprintln!("错误: {}", e);
        std::process::exit(1);
    }
    
    if config.show_count {
        writeln!(output).unwrap();
        writeln!(output, "{} 个目录, {} 个文件", stats.dirs, stats.files).unwrap();
    }
    
    // 输出到文件或标准输出
    if let Some(output_file) = args.output {
        match File::create(&output_file) {
            Ok(mut file) => {
                if let Err(e) = file.write_all(output.as_bytes()) {
                    eprintln!("写入文件 '{}' 时出错: {}", output_file, e);
                    std::process::exit(1);
                }
                println!("树状结构已输出到文件 '{}'.", output_file);
            }
            Err(e) => {
                eprintln!("创建文件 '{}' 时出错: {}", output_file, e);
                std::process::exit(1);
            }
        }
    } else {
        print!("{}", output);
    }
}

fn traverse_directory(
    dir: &Path,
    config: &TreeConfig,
    prefix: &str,
    _is_last: bool,
    depth: usize,
    stats: &mut Stats,
    output: &mut String,
) -> io::Result<()> {
    // 检查深度限制
    if let Some(max_depth) = config.max_depth {
        if depth >= max_depth {
            return Ok(());
        }
    }
    
    // 读取目录条目
    let mut entries: Vec<_> = fs::read_dir(dir)?
        .filter_map(|entry| entry.ok())
        .collect();
    
    // 根据配置过滤条目
    entries.retain(|entry| {
        let file_name = entry.file_name();
        let name = file_name.to_string_lossy();
        
        // 除非指定 --all 选项，否则跳过隐藏文件
        if !config.show_all && name.starts_with('.') {
            return false;
        }
        
        // 如果 dirs_only 为 true，仅显示目录
        if config.dirs_only {
            if let Ok(metadata) = entry.metadata() {
                return metadata.is_dir();
            }
        }
        
        true
    });
    
    // 排序条目：目录在前，然后是文件，都按字母顺序排列
    entries.sort_by(|a, b| {
        let a_is_dir = a.metadata().map(|m| m.is_dir()).unwrap_or(false);
        let b_is_dir = b.metadata().map(|m| m.is_dir()).unwrap_or(false);
        
        match (a_is_dir, b_is_dir) {
            (true, false) => std::cmp::Ordering::Less,
            (false, true) => std::cmp::Ordering::Greater,
            _ => a.file_name().cmp(&b.file_name()),
        }
    });
    
    for (i, entry) in entries.iter().enumerate() {
        let is_last_entry = i == entries.len() - 1;
        let path = entry.path();
        let file_name = entry.file_name();
        let name = file_name.to_string_lossy();
        
        // 确定树状符号
        let current_symbol = if is_last_entry { "└── " } else { "├── " };
        let next_prefix = if is_last_entry {
            format!("{}    ", prefix)
        } else {
            format!("{}│   ", prefix)
        };
        
        // 打印当前条目（带可选的文件大小）
        let entry_line = if config.show_size {
            if let Ok(metadata) = entry.metadata() {
                if metadata.is_dir() {
                    format!("{}{}{}/", prefix, current_symbol, name)
                } else {
                    let size = format_file_size(metadata.len());
                    format!("{}{}{} [{}]", prefix, current_symbol, name, size)
                }
            } else {
                format!("{}{}{}", prefix, current_symbol, name)
            }
        } else {
            format!("{}{}{}", prefix, current_symbol, name)
        };
        
        writeln!(output, "{}", entry_line).unwrap();
        
        // 更新统计信息，如果是目录则递归遍历
        if let Ok(metadata) = entry.metadata() {
            if metadata.is_dir() {
                stats.dirs += 1;
                if let Err(e) = traverse_directory(&path, config, &next_prefix, is_last_entry, depth + 1, stats, output) {
                    eprintln!("警告: 无法访问目录 '{}': {}", path.display(), e);
                }
            } else {
                stats.files += 1;
            }
        }
    }
    
    Ok(())
}

fn format_file_size(size: u64) -> String {
    const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
    let mut size = size as f64;
    let mut unit_index = 0;
    
    while size >= 1024.0 && unit_index < UNITS.len() - 1 {
        size /= 1024.0;
        unit_index += 1;
    }
    
    if unit_index == 0 {
        format!("{} {}", size as u64, UNITS[unit_index])
    } else {
        format!("{:.1} {}", size, UNITS[unit_index])
    }
}
